#include "nominallabelextract.h"

#include <errno.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>

static double
qfunc_data (Data* data)
{
  return qfunc (data->alpha,
                data->n_workers,
                data->n_labels,
                data->logbeta,
                data->n_images,
                data->pz,
                data->priorz,
                data->ratings,
                data->n_ratings,
                &data->prioralpha.base,
                &data->priorlogbeta.base,
                data->symmetric);
}

static void
estep_data (Data* data)
{
  return estep (data->alpha,
                data->n_labels,
                data->logbeta,
                data->n_images,
                data->pz,
                data->priorz,
                data->ratings,
                data->n_ratings,
                data->symmetric);
}

static double
mstep_data (Data* data)
{
  return mstep (data->alpha,
                data->n_workers,
                data->n_labels,
                data->logbeta,
                data->n_images,
                data->pz,
                data->priorz,
                data->ratings,
                data->n_ratings,
                &data->prioralpha.base,
                &data->priorlogbeta.base,
                data->symmetric,
                1e-3);
}

static void
em (Data* data)
{
  double q = qfunc_data (data);
  double lastq = q;
  const double epsilon = 1e-5;

  do
    {
      lastq = q;
      estep_data (data);
      mstep_data (data);
      q = qfunc_data (data);
    }
  while (fabs ((q - lastq) / lastq) > epsilon);
}

static void
print_alpha (FILE*       fp,
             const Data* data)
{
  for (unsigned int i = 0; i < data->n_workers; ++i)
    {
      fprintf (fp, "Alpha[%u] =", i);

      for (unsigned int j = 0; j < data->n_labels; ++j)
        {
          for (unsigned int k = 0; k < data->n_labels; ++k)
            {
              if (k == j)
                {
                  /* do nothing */
                }
              else if (data->symmetric && k < j) 
                {
                  /* do nothing */
                }
              else
                {
                  fprintf (fp, 
                           " %f", 
                           data->alpha[i * data->n_labels * data->n_labels +
                                       j * data->n_labels +
                                       k]);
                }
            }
        }

      fprintf (fp, "\n");
    }
}

static void
print_logbeta (FILE*       fp,
               const Data* data)
{
  for (unsigned int i = 0; i < data->n_images; ++i)
    {
      fprintf (fp, "Beta[%u] = %f\n", i, exp (data->logbeta[i]));
    }
}

static void
print_pz (FILE*         fp,
          const Data*   data)
{
  for (unsigned int i = 0; i < data->n_images; ++i)
    {
      fprintf (fp, "P(Z(%u)=(k>0)) =", i);

      for (unsigned int j = 1; j < data->n_labels; ++j)
        {
          fprintf (fp, " %f", data->pz[j + i * data->n_labels]);
        }

      fprintf (fp, "\n");
    }
}

static void
print_parameters (FILE*         fp,
                  const Data*   data)
{
  print_alpha (fp, data);
  print_logbeta (fp, data);
  print_pz (fp, data);
}

static void
usage (void)
{
  fprintf (stderr, 
           "usage: nominallabelextract [-h|-s] [datafile]\n"
           "\n"
           "  datafile: specify file to read input data from.\n"
           "      if datafile eq '-' use standard input\n"
           "      if not specified, \n"
           "        if standard input is not a tty, use standard input\n"
           "        if standard input is a tty, display this message.\n"
           "\n"
           "  -h: display this message.\n"
           "\n"
           "  -s: use symmetric error model.  mostly this was to check my implementation\n"
           "      against Whitehill et. al., so you probably don't want this.\n"
          );

  exit (1);
}

int 
main (int       argc, 
      char*     argv[])
{
  Data data;
  int symmetric = 0;

  if (argc > 1 && strcmp (argv[1], "-h") == 0)
    {
      usage ();
    }
  else if (argc > 1 && strcmp (argv[1], "-s") == 0)
    {
      symmetric = 1;
      ++argv;
      --argc;
    }

  if (argc == 1 && isatty (0))
    {
      usage ();
    }

  FILE* fp = (argc == 1 || strcmp (argv[1], "-") == 0) ? stdin
                                                       : fopen (argv[1], "r");

  if (fp == NULL)
    {
      fprintf (stderr, 
               "can't open %s: %s\n",
               argv[1],
               strerror (errno));
      exit (1);
    }

  if (read_data (fp, &data, symmetric) == 0)
    {
      em (&data);
      print_parameters (stdout, &data);
    }

  fclose (fp);
  data_destruct (&data);

  return 0;
}
